/*
 * Description: performance hooking aspect.
 */
#include "pch.h"
#include <functional>
#include <algorithm>
#include <Windows.h>
#include <thread>
#include <TlHelp32.h>
#include <sddl.h>
#include <shlobj.h>
#include <unordered_set>
#include "aspect.h"
#include "messenger.h"
#include "include/detours.h"
#include "performance/perf_constants.h"
#include "performance/perf_scene.h"

namespace PERF {

POriginalCreateProcessA g_pOriginalCreateProcessA;
POriginalCreateProcessW g_pOriginalCreateProcessW;
pfShellExecuteExW g_pOriginalShellExecuteExW;
std::string WCharToMByteString(LPCWSTR lpcwszStr);

const std::string MSG_SEP = "#";
const std::string UNUSE_CMD_LINE = "UNUSECMD";
const std::string AUTOCAD_FILE_NAME_EXTENSION = ".dwg";
const std::string FILE_OPEN_ACTION = "open";
const std::string FILE_RUNAS_ACTION = "runas";
static const UINT MAX_CMD_LEN = 32;
static const int INVOKECOMMAND_FUN_IDX = 4;  // 虚函数位置
static const UINT QUERY_FILE_IDX = 0xFFFFFFFF;  // 查询文件索引,0xFFFFFFFF，则 DragQueryFile 将返回ALL文件计数
const std::unordered_set<std::string> PROCESS_FILTER_SET = {};

typedef HRESULT(__stdcall *INVOKECOMMAND)(IContextMenu *pcm, CMINVOKECOMMANDINFO *pici);
ULONG_PTR *gVtbl = nullptr;                 // 虚函数表地址
ULONG_PTR *gAddr = nullptr;                 // 虚函数地址指针
INVOKECOMMAND pOldInvokeCommand = nullptr;  // 虚函数地址

std::string GetTargetFileExtensionPath(HDROP hDrop)
{
    // 获取操作中被选中的文件的数目
    UINT n = DragQueryFile(hDrop, QUERY_FILE_IDX, NULL, 0);
    for (UINT i = 0; i < n; i++) {
        // 获取被选中的文件路径
        char szPath[MAX_PATH] = { 0 };
        DragQueryFile(hDrop, i, szPath, MAX_PATH);
        std::string szDest = szPath;
        int len = AUTOCAD_FILE_NAME_EXTENSION.length();
        if (szDest.length() >= len) {
            std::string last_four = szDest.substr(szDest.size() - len);
            if (last_four == AUTOCAD_FILE_NAME_EXTENSION) {
                return szDest;
            }
        }
    }
    return "";
}

void CheckDragFileList(IContextMenu *pcm)
{
    IObjectWithSelection *pSelection = nullptr;
    HRESULT hr = pcm->QueryInterface(IID_IObjectWithSelection, (VOID **)&pSelection);
    if (!SUCCEEDED(hr)) {
        return;
    }
    IDataObject *pObject = nullptr;
    hr = pSelection->GetSelection(IID_IDataObject, (VOID **)&pObject);
    if (!SUCCEEDED(hr)) {
        pSelection->Release();
        return;
    }
    STGMEDIUM stg = { TYMED_HGLOBAL };
    FORMATETC etc = { CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL };
    // 获取文件列表指针,它们存储在HDROP格式中
    hr = pObject->GetData(&etc, &stg);
    if (SUCCEEDED(hr)) {
        // 取得HDROP句柄
        HDROP hDrop = (HDROP)GlobalLock(stg.hGlobal);
        if (!hDrop) {
            pObject->Release();
            pSelection->Release();
            return;
        }
        std::string szPath = GetTargetFileExtensionPath(hDrop);
        if (!szPath.empty()) {
            std::thread tSendMessageToService(SendMouseDragMessageToService, PNAME_AUTOCAD, szPath);
            tSendMessageToService.detach();
        }
        GlobalUnlock(stg.hGlobal);
        ReleaseStgMedium(&stg);
    }
    pObject->Release();
    pSelection->Release();
    return;
}

HRESULT g_pOriginalInvokeCommand(IContextMenu *pcm, CMINVOKECOMMANDINFO *pici)
{
    HRESULT hr = S_OK;
    char cmd[MAX_CMD_LEN] = { 0 };
    do {
        if (!pcm || !pici) {
            break;
        }
        // 获取操作的动作
        if (!IS_INTRESOURCE(pici->lpVerb)) {
            break;
        }
        hr = pcm->GetCommandString((UINT_PTR)pici->lpVerb, GCS_VERBA, NULL, cmd, MAX_CMD_LEN);
        if (FAILED(hr)) {
            break;
        }
        // 判断是不是运行操作
        if (_stricmp(cmd, FILE_OPEN_ACTION.c_str()) && _stricmp(cmd, FILE_RUNAS_ACTION.c_str())) {
            break;
        }
        // 如果匹配到目标文件就直接返回 选中多个目标的情况下只命中第1个
        CheckDragFileList(pcm);
    } while (0);
    // 调用原始函数
    if (pOldInvokeCommand) {
        hr = pOldInvokeCommand(pcm, pici);
    }
    return hr;
}

void HookInvokeCommandVirtualTable()
{
    if (pOldInvokeCommand) {
        return;
    }
    IContextMenu *pcm = NULL;
    HRESULT hr = CDefFolderMenu_Create2(NULL, NULL, 0, NULL, NULL, NULL, 0, NULL, &pcm);
    if (FAILED(hr)) {
        return;
    }
    gVtbl = *(ULONG_PTR **)pcm;
    gAddr = (ULONG_PTR *)&(gVtbl[INVOKECOMMAND_FUN_IDX]);
    pOldInvokeCommand = (INVOKECOMMAND)*gAddr;
}

LONG AppendAspect()
{
    DetourTransactionBegin();
    HookInvokeCommandVirtualTable();
    DetourUpdateThread(GetCurrentThread());
    g_pOriginalCreateProcessA = (POriginalCreateProcessA)DetourFindFunction("kernel32.dll", "CreateProcessA");
    g_pOriginalCreateProcessW = (POriginalCreateProcessW)DetourFindFunction("kernel32.dll", "CreateProcessW");
    g_pOriginalShellExecuteExW = ShellExecuteExW;
    DetourAttach((PVOID *)&g_pOriginalCreateProcessA, AspectedCreateProcessA);
    DetourAttach((PVOID *)&g_pOriginalCreateProcessW, AspectedCreateProcessW);
    DetourAttach((PVOID *)&g_pOriginalShellExecuteExW, AspectedShellExecuteExW);
    DetourAttach(&(PVOID &)pOldInvokeCommand, g_pOriginalInvokeCommand);
    LONG ret = DetourTransactionCommit();
    return ret;
}

LONG RemoveAspect()
{
    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());
    DetourDetach((PVOID *)&g_pOriginalCreateProcessA, AspectedCreateProcessA);
    DetourDetach((PVOID *)&g_pOriginalCreateProcessW, AspectedCreateProcessW);
    DetourDetach((PVOID *)&g_pOriginalShellExecuteExW, AspectedShellExecuteExW);
    DetourDetach((PVOID *)&pOldInvokeCommand, g_pOriginalInvokeCommand);
    LONG ret = DetourTransactionCommit();
    pOldInvokeCommand = nullptr;
    return ret;
}

BOOL WINAPI AspectedCreateProcessA(_In_opt_ LPCSTR lpApplicationName, _Inout_opt_ LPCSTR lpCommandLine,
                                   _In_opt_ LPSECURITY_ATTRIBUTES lpProcessAttributes,
                                   _In_opt_ LPSECURITY_ATTRIBUTES lpThreadAttributes, _In_ BOOL bInheritHandles,
                                   _In_ DWORD dwCreationFlags, _In_opt_ LPVOID lpEnvironment,
                                   _In_opt_ LPCSTR lpCurrentDirectory, _In_ LPSTARTUPINFOW lpStartupInfo,
                                   _Out_ LPPROCESS_INFORMATION lpProcessInformation)
{
    for (std::string filteredProcessName : PROCESS_FILTER_SET) {
        if (std::string(lpApplicationName).find(filteredProcessName) != std::string::npos) {
            if (GetProcessIdByName(filteredProcessName.c_str()) != 0) {
                std::thread tSendMessageToService(SendMessageToServiceA, lpApplicationName, -1, lpCommandLine);
                tSendMessageToService.detach();
                return true;
            }
        }
    }
    BOOL ret = g_pOriginalCreateProcessA(lpApplicationName, lpCommandLine, lpProcessAttributes, lpThreadAttributes,
                                         bInheritHandles, dwCreationFlags, lpEnvironment, lpCurrentDirectory,
                                         lpStartupInfo, lpProcessInformation);

    if (lpProcessInformation == nullptr || lpApplicationName == nullptr || lpCommandLine == nullptr) {
        return ret;
    }
    std::string strApplicationName = lpApplicationName;
    std::string strCommandLine = lpCommandLine;
    DWORD dwPid = lpProcessInformation->dwProcessId;
    std::thread tSendMessageToService(SendMessageToServiceA, strApplicationName, dwPid, strCommandLine);
    tSendMessageToService.detach();

    return ret;
}

BOOL WINAPI AspectedCreateProcessW(_In_opt_ LPCWSTR lpApplicationName, _Inout_opt_ LPCWSTR lpCommandLine,
                                   _In_opt_ LPSECURITY_ATTRIBUTES lpProcessAttributes,
                                   _In_opt_ LPSECURITY_ATTRIBUTES lpThreadAttributes, _In_ BOOL bInheritHandles,
                                   _In_ DWORD dwCreationFlags, _In_opt_ LPVOID lpEnvironment,
                                   _In_opt_ LPCWSTR lpCurrentDirectory, _In_ LPSTARTUPINFOW lpStartupInfo,
                                   _Out_ LPPROCESS_INFORMATION lpProcessInformation)
{
    for (std::string filteredProcessName : PROCESS_FILTER_SET) {
        if (WCharToMByteString(lpApplicationName).find(filteredProcessName) != std::string::npos) {
            if (GetProcessIdByName(filteredProcessName.c_str()) != 0) {
                std::thread tSendMessageToService(SendMessageToServiceW, lpApplicationName, -1, lpCommandLine);
                tSendMessageToService.detach();
                return true;
            }
        }
    }
    BOOL ret = g_pOriginalCreateProcessW(lpApplicationName, lpCommandLine, lpProcessAttributes, lpThreadAttributes,
                                         bInheritHandles, dwCreationFlags, lpEnvironment, lpCurrentDirectory,
                                         lpStartupInfo, lpProcessInformation);

    if (lpProcessInformation == nullptr || lpApplicationName == nullptr || lpCommandLine == nullptr) {
        return ret;
    }
    std::wstring wstrApplicationName = lpApplicationName;
    std::wstring wstrCommandLine = lpCommandLine;
    DWORD dwPid = lpProcessInformation->dwProcessId;
    std::thread tSendMessageToService(SendMessageToServiceW, wstrApplicationName, dwPid, wstrCommandLine);
    tSendMessageToService.detach();

    return ret;
}

void SendMessage(const std::string &appPath, const std::string &cmdline, DWORD dwpid)
{
    if (appPath.size() == 0 || cmdline.size() == 0) {
        return;
    }
    std::string appName = appPath.substr(appPath.find_last_of('\\') + 1);
    std::transform(appName.begin(), appName.end(), appName.begin(), ::tolower);
    if (find(INEFFECTIVE_APPS.begin(), INEFFECTIVE_APPS.end(), appName) != INEFFECTIVE_APPS.end()) {
        return;
    }
    Messenger messenger;
    messenger.Send(appInit + MSG_SEP + appPath + MSG_SEP + appName + MSG_SEP + std::to_string(dwpid) + MSG_SEP +
                   cmdline);
}

void SendMessageToServiceA(std::string strApplicationName, DWORD dwpid, std::string strCmdLine)
{
    SendMessage(strApplicationName, strCmdLine, dwpid);
}

void SendMessageToServiceW(std::wstring wstrApplicationName, DWORD dwpid, std::wstring wstrCmdLine)
{
    std::string appPath = WCharToMByteString(wstrApplicationName.c_str());
    std::string cmdline = WCharToMByteString(wstrCmdLine.c_str());
    SendMessage(appPath, cmdline, dwpid);
}

std::string WCharToMByteString(LPCWSTR lpcwszStr)
{
    if (lpcwszStr == nullptr) {
        return "";
    }
    std::string str;
    DWORD dwMinSize = 0;
    LPSTR lpszStr = nullptr;
    dwMinSize = WideCharToMultiByte(CP_OEMCP, 0, lpcwszStr, -1, nullptr, 0, nullptr, FALSE);
    if (0 == dwMinSize) {
        return "";
    }
    lpszStr = new char[dwMinSize];
    WideCharToMultiByte(CP_OEMCP, 0, lpcwszStr, -1, lpszStr, dwMinSize, nullptr, FALSE);
    str = lpszStr;
    delete[] lpszStr;
    return str;
}

BOOL WINAPI AspectedShellExecuteExW(SHELLEXECUTEINFOW *pExecInfo)
{
    std::thread tSendMessageToService(SendMessageToService, PNAME_EXPLORER);
    tSendMessageToService.detach();
    return g_pOriginalShellExecuteExW(pExecInfo);
}

void SendMessageToService(const std::string &exeName)
{
    DWORD currentPid = GetProcessIdByName(exeName.c_str());
    std::string exePath = GetProcessPathByPid(currentPid);
    // Notify vmagent
    Messenger messenger;
    messenger.Send(vmBoot + MSG_SEP + exePath + MSG_SEP + exeName + MSG_SEP + std::to_string(currentPid) + MSG_SEP +
                   UNUSE_CMD_LINE);
}

void SendMouseDragMessageToService(const std::string &exeName, const std::string &exePath)
{
    DWORD currentPid = GetProcessIdByName(exeName.c_str());
    Messenger messenger;
    messenger.Send(fileOpen + MSG_SEP + exePath + MSG_SEP + exeName + MSG_SEP + std::to_string(currentPid) + MSG_SEP +
                   UNUSE_CMD_LINE);
}

DWORD GetProcessIdByName(LPCTSTR exeName)
{
    DWORD nPid = 0;
    if (exeName == nullptr) {
        return nPid;
    }
    PROCESSENTRY32 pe = { 0 };
    pe.dwSize = sizeof(PROCESSENTRY32);
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
    if (hSnapshot != INVALID_HANDLE_VALUE && Process32First(hSnapshot, &pe)) {
        do {
            if (pe.szExeFile != nullptr && lstrcmpi(exeName, pe.szExeFile) == 0) {
                nPid = pe.th32ProcessID;
                break;
            }
        } while (Process32Next(hSnapshot, &pe));
    }
    if (hSnapshot != INVALID_HANDLE_VALUE && hSnapshot != nullptr) {
        CloseHandle(hSnapshot);
    }
    return nPid;
}

std::string GetProcessPathByPid(DWORD processId)
{
    HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId);
    if (hProcess == nullptr || hProcess == INVALID_HANDLE_VALUE) {
        return "";
    }
    WCHAR buffer[MAX_PATH] = { 0 };
    DWORD size = MAX_PATH;
    if (QueryFullProcessImageNameW(hProcess, 0, buffer, &size)) {
        int bufferSize = WideCharToMultiByte(CP_UTF8, 0, buffer, -1, NULL, 0, NULL, NULL);
        if (bufferSize == 0) {
            CloseHandle(hProcess);
            return "";
        }
        std::string exePath(bufferSize, 0);
        if (!WideCharToMultiByte(CP_UTF8, 0, buffer, -1, &exePath[0], bufferSize, NULL, NULL)) {
            CloseHandle(hProcess);
            return "";
        }
        exePath.resize(bufferSize - 1);
        CloseHandle(hProcess);
        return exePath;
    }
    CloseHandle(hProcess);
    return "";
}
}